'''
DeepSort跟踪算法,来自https://github.com/mikel-brostrom/Yolov5_DeepSort_OSNet
xiaohe
2022/05/27
'''
import os
import numpy as np

from sort.tracker import Tracker
from sort.detection import Detection
from sort.np_matching import NearestNeighborDistanceMetric
from sort.roi_feature_extractor import feature_extracor_network

__all__ = ['DeepSort']

class DeepSort(object):
    def __init__(self,  max_dist=0.2, max_iou_distance=0.7, max_age=70, n_init=3, nn_budget=100, **kwargs):
        """
        DeepSort追踪算法
        :param model_path: 模型地址
        :param max_dist: The matching threshold. Samples with larger distance are considered an invalid match
        :param max_iou_distance: Gating threshold. Associations with cost larger than this value are disregarded.
        :param MAX_AGE: Maximum number of missed misses before a track is deleted
        :param N_INIT: Number of frames that a track remains in initialization phase
        :param NN_BUDGET: Maximum size of the appearance descriptors gallery
        """
        self.extractor = feature_extracor_network(**kwargs)
        metric = NearestNeighborDistanceMetric("cosine", max_dist, nn_budget)
        self.tracker = Tracker(metric, max_iou_distance=max_iou_distance, max_age=max_age, n_init=n_init)

    def update(self, bbox_xyxy, confidences, classes, ori_img):
        self.height, self.width = ori_img.shape[:2]
        # generate detections
        features = self._get_features(bbox_xyxy, ori_img)
        bbox_tlwh = self._xyxy_to_tlwh(bbox_xyxy)
        detections = [Detection(bbox_tlwh[i], conf, features[i]) for i, conf in enumerate(confidences)]

        # update tracker
        self.tracker.predict()
        self.tracker.update(detections, classes, confidences)

        # output bbox identities
        outputs = []
        for track in self.tracker.tracks:
            if not track.is_confirmed() or track.time_since_update > 1:
                continue
            x1, y1, x2, y2 = track.to_x1y1x2y2()
            outputs.append(np.array([x1, y1, x2, y2, track.track_id, track.class_id, track.conf]))
        if len(outputs) > 0:
            outputs = np.stack(outputs, axis=0)
        return outputs

    def predict(self):
        self.tracker.predict()
        # output bbox identities
        outputs = []
        for track in self.tracker.tracks:
            if not track.is_confirmed() or track.time_since_update > 1:
                continue
            x1, y1, x2, y2 = track.to_x1y1x2y2()
            outputs.append(np.array([x1, y1, x2, y2, track.track_id, track.class_id, track.conf]))
        if len(outputs) > 0:
            outputs = np.stack(outputs, axis=0)
        return outputs

    def increment_ages(self):
        self.tracker.increment_ages()

    def _xyxy_to_tlwh(self, bbox_xyxy):
        bbox_tlwh = bbox_xyxy.copy()
        bbox_tlwh[:, 2] = bbox_xyxy[:, 2] - bbox_xyxy[:, 0]
        bbox_tlwh[:, 3] = bbox_xyxy[:, 3] - bbox_xyxy[:, 1]
        return bbox_tlwh

    def _get_features(self, bboxes, ori_img):
        im_crops = []
        for box in bboxes:
            x1, y1, x2, y2 = [int(i) for i in box]
            im = ori_img[y1:y2, x1:x2]
            im_crops.append(im)
        if im_crops:
            features = self.extractor(im_crops)
        else:
            features = np.array([])
        return features